{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "import os\n",
    "from glob import glob\n",
    "import re\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.backends.backend_pdf import PdfPages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/media/liamcli/fastfiles/results/icml2020'\n",
    "plt.style.use('ggplot')\n",
    "params = {'legend.fontsize': 14}\n",
    "plt.rcParams.update(params)\n",
    "plt.rcParams.update({'font.size': 12})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_arch_params(folder, softmax=False):\n",
    "  files = glob(os.path.join(folder, 'one_shot_architecture*.obj'))\n",
    "  ordered = []\n",
    "  for f in files:\n",
    "    start = f.find('architecture_') + len('architecture_')\n",
    "    end = f.find('.obj') \n",
    "    epoch = int(f[start:end])\n",
    "    ordered.append((epoch, f))\n",
    "  ordered = sorted(ordered, key=lambda x:x[0])\n",
    "  archs = []\n",
    "  for e, f in ordered:\n",
    "    with open(f, 'rb') as b:\n",
    "      arch = pickle.load(b)\n",
    "      if softmax:\n",
    "        for i, p in enumerate(arch):\n",
    "          normalized = np.exp(p)\n",
    "          arch[i] = normalized / np.sum(normalized, axis=-1, keepdims=True)\n",
    "          \n",
    "      archs.append(arch)\n",
    "  return archs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def entropy(p, axis):\n",
    "  return -np.sum(p*np.log(p), axis=axis)\n",
    "\n",
    "def get_entropy(history):\n",
    "  result = {\n",
    "    'operations': [],\n",
    "    'edges': [],\n",
    "  }\n",
    "  for params in history:\n",
    "    result['operations'].append(np.average(entropy(params[0],axis=-1)))\n",
    "    result['edges'].append(np.average([entropy(p, axis=-1) for p in params[1:]]))\n",
    "  return result\n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = {\n",
    "  'pc_darts': [\n",
    "    'ameet/search_space_1/search-pc_darts_space1_seed3500', \n",
    "    'ameet/search_space_2/search-pc_darts_space2_seed3501',\n",
    "    'ameet/search_space_3/search-pc_darts_space3_seed3502'\n",
    "  ],\n",
    "  'epc_edarts':[\n",
    "    'ameet/search_space_1/search-epc_edarts_space1_seed5100', \n",
    "    'ameet/search_space_2/search-epc_edarts_space2_seed5101',\n",
    "    'ameet/search_space_3/search-epc_edarts_space3_seed5102'\n",
    "  ]\n",
    "}\n",
    "styles = {\n",
    "  'pc_darts': '-',\n",
    "  'epc_edarts': '--'\n",
    "}\n",
    "colors = {\n",
    "  'pc_darts': 'b',\n",
    "  'epc_edarts': 'r'\n",
    "}\n",
    "labels = {\n",
    "  'pc_darts': 'PC-DARTS',\n",
    "  'epc_edarts': 'GAEA PC-DARTS',\n",
    "  'darts': 'DARTS (first order)',\n",
    "  'edarts': 'GAEA DARTS (bi-level)'\n",
    "}\n",
    "pdf = PdfPages('1shot1_entropy.pdf')\n",
    "for s in [1, 2, 3]:\n",
    "  for method in ['pc_darts', 'epc_edarts']:\n",
    "    exp_dir = os.path.join(root, folders[method][s-1])\n",
    "    data = get_arch_params(os.path.join(root, exp_dir+'*'), True if method=='pc_darts' else False)\n",
    "    ent = get_entropy(data)\n",
    "    labeled=False\n",
    "    for f in ['operations']:\n",
    "      if not labeled:\n",
    "        plt.plot(ent[f], label=labels[method], linestyle=styles[method], color=colors[method])\n",
    "        labeled=True\n",
    "      else:\n",
    "        plt.plot(ent[f], linestyle=styles[method], color=colors[method])\n",
    "  plt.title('NAS-Bench-1Shot1: Space {} Entropy'.format(s))\n",
    "  plt.xlim([0, 50])\n",
    "  plt.xlabel('Epochs')\n",
    "  plt.ylabel('Entropy')\n",
    "  plt.legend()\n",
    "  plt.tight_layout()\n",
    "  pdf.savefig()\n",
    "  plt.close()\n",
    "pdf.close()\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_history_entropy(folder, softmax=False):\n",
    "  data = pickle.load(open(os.path.join(folder, 'history.pkl'), 'rb'))\n",
    "  results = {\n",
    "    'operations': [],\n",
    "    'edges': []\n",
    "  }\n",
    "  for i, r in enumerate(data['alphas']):\n",
    "    es = []\n",
    "    for ct in r:\n",
    "      prob = r[ct]\n",
    "      if softmax:\n",
    "        prob = np.exp(prob)\n",
    "        prob = prob / np.sum(prob, axis=-1, keepdims=True)\n",
    "      es.append(np.average(entropy(prob, axis=-1)))\n",
    "    results['operations'].append(es)\n",
    "  #if 'edges' in data:\n",
    "  #  for i, r in enumerate(data['edges']):\n",
    "  #    es = []\n",
    "  #    for ct in r:\n",
    "  #      offset = 0\n",
    "  #      n_inputs = 2\n",
    "  #      for n in range(4):\n",
    "  #        es.extend(entropy(r[ct][offset:offset+n_inputs], axis=-1))\n",
    "  #        offset += n_inputs\n",
    "  #        n_inputs += 1\n",
    "  #    results['edges'].append(np.average(es))\n",
    "  return results\n",
    "      \n",
    "  \n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = {\n",
    "  'epc_edarts': [\n",
    "    'ameet/search-pcdarts-eedarts-cifar10-3103'\n",
    "  ],\n",
    "  'pc_darts':[\n",
    "    'nina/search-pcdarts-darts-cifar10-3600'\n",
    "  ]\n",
    "}\n",
    "styles = {\n",
    "  'pc_darts': '-',\n",
    "  'epc_edarts': '--'\n",
    "}\n",
    "colors = {\n",
    "  'pc_darts': 'b',\n",
    "  'epc_edarts': 'r'\n",
    "}\n",
    "pdf = PdfPages('darts_entropy.pdf')\n",
    "for method in ['pc_darts', 'epc_edarts']:\n",
    "  exp_dir = os.path.join(root, folders[method][0])\n",
    "  #data = get_arch_params(os.path.join(root, exp_dir+'*'), True if method=='pc_darts' else False)\n",
    "  ent = get_history_entropy(os.path.join(root, exp_dir), True if method=='pc_darts' else False)\n",
    "  labeled=False\n",
    "  for f in ent:\n",
    "    if not labeled:\n",
    "      plt.plot([t[0] for t in ent[f]], label=labels[method], linestyle=styles[method], color=colors[method])\n",
    "      #plt.plot([t[1] for t in ent[f]], linestyle=styles[method], color=colors[method])\n",
    "      labeled=True\n",
    "    else:\n",
    "      plt.plot(ent[f], linestyle=styles[method], color=colors[method])\n",
    "plt.title('DARTS Search Space Entropy'.format(s))\n",
    "plt.xlim([0, 50])\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Entropy')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "pdf.savefig()\n",
    "plt.close()\n",
    "pdf.close()\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/media/liamcli/fastfiles/results/icml2020'\n",
    "folders = {\n",
    "  'darts': [\n",
    "    'ameet/search-nas-bench-201-fdarts-cifar10-4400'\n",
    "  ],\n",
    "  'edarts':[\n",
    "    'search-nas-bench-201-edarts-cifar10-6500'\n",
    "  ]\n",
    "}\n",
    "styles = {\n",
    "  'darts': '-',\n",
    "  'edarts': '--'\n",
    "}\n",
    "colors = {\n",
    "  'darts': 'b',\n",
    "  'edarts': 'r'\n",
    "}\n",
    "pdf = PdfPages('nasbench201_entropy.pdf')\n",
    "for method in ['darts', 'edarts']:\n",
    "  exp_dir = os.path.join(root, folders[method][0])\n",
    "  #data = get_arch_params(os.path.join(root, exp_dir+'*'), True if method=='pc_darts' else False)\n",
    "  ent = get_history_entropy(os.path.join(root, exp_dir), True if method=='darts' else False)\n",
    "  labeled=False\n",
    "  for f in ent:\n",
    "    if not labeled:\n",
    "      plt.plot([t[0] for t in ent[f]], label=labels[method], linestyle=styles[method], color=colors[method])\n",
    "      #plt.plot([t[1] for t in ent[f]], linestyle=styles[method], color=colors[method])\n",
    "      labeled=True\n",
    "    else:\n",
    "      plt.plot(ent[f], linestyle=styles[method], color=colors[method])\n",
    "plt.title('NAS-Bench-201:  CIFAR-10 Entropy'.format(s))\n",
    "plt.xlim([0, 50])\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Entropy')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "pdf.savefig()\n",
    "plt.close()\n",
    "pdf.close()\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
